# -*-coding:utf-8-*-
import numpy as np
import cv2


def show_result_pyplot(model,
                       img,
                       result,
                       score_thr=0.7,
                       title='result',
                       wait_time=0,
                       out_file=None):
    """Visualize the detection results on the image.

    Args:
        model (nn.Module): The loaded detector.
        img (str or np.ndarray): Image filename or loaded image.
        result (tuple[list] or list): The detection result, can be either
            (bbox, segm) or just bbox.
        score_thr (float): The threshold to visualize the bboxes and masks.
        title (str): Title of the pyplot figure.
        wait_time (float): Value of waitKey param.
                Default: 0.
    """
    if hasattr(model, 'module'):
        model = model.module
    return model.show_result(
        img,
        result,
        score_thr=score_thr,
        show=False,
        wait_time=wait_time,
        win_name=title,
        bbox_color=(72, 101, 241),
        text_color=(72, 101, 241),
        out_file=out_file)


class SegDetect:
    def __init__(self, img):
        self.img = img
        self.index, self.bitnot1, self.bitnot2 = None, None, None

    def detect_seg(self, bbox_result, segment_func, model):
        #  extract bbox for segmentation
        image = cv2.imread(self.img)
        bboxes = np.concatenate((bbox_result[0], bbox_result[1]), axis=0)
        bboxes = self.filtering_boxes(bboxes, 0.7)  # 得到每张图片的bboxes -> list[list]
        if not self.index:
            _, _, self.index, self.bitnot1, self.bitnot2 = segment_func(image)
        for i, bbox in enumerate(bboxes):
            local = image[bbox[1]:bbox[3], bbox[0]:bbox[2], :]
            _, local_result, _, _, _ = segment_func(local, self.index, self.bitnot1, self.bitnot2)
            image[bbox[1]:bbox[3], bbox[0]:bbox[2], :] = local_result
            # cv_show('name', detect_img)
        return show_result_pyplot(model, image, bbox_result), image

    @staticmethod
    def filtering_boxes(bboxes, score_thr):
        scores = bboxes[:, -1]
        inds = scores > score_thr
        return np.round(bboxes[inds]).astype('int32').tolist()


if __name__ == '__main__':
    pass
